package com.pw.study.flink.watermark

import com.pw.study.flink.source.{WordBean, WordSource}
import org.apache.flink.streaming.api.TimeCharacteristic
import org.apache.flink.streaming.api.functions.{AssignerWithPeriodicWatermarks, AssignerWithPunctuatedWatermarks}
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.function.ProcessWindowFunction
import org.apache.flink.streaming.api.watermark.Watermark
import org.apache.flink.streaming.api.windowing.time.Time
import org.apache.flink.streaming.api.windowing.windows.TimeWindow
import org.apache.flink.util.Collector

/**
 * @Desc: TODO
 * @author: future
 * @since: 2022/5/15-20:48
 */
object CustomWaterMarkPunctuated {
  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(2)
    env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime)
    //设置周期性watermark的时间间隔
    env.getConfig.setAutoWatermarkInterval(3000L)

    val ds: DataStream[WordBean] = env.addSource(new WordSource)
    val dsWatermark: DataStream[WordBean] = ds.assignTimestampsAndWatermarks(
     new AssignerWithPunctuatedWatermarks[WordBean] {
       override def checkAndGetNextWatermark(t: WordBean, ts: Long): Watermark = {
         println("checkAndGetNextWatermark..........")
         new Watermark(ts)
       }

       override def extractTimestamp(element: WordBean, ts: Long): Long = {
         println("extractTimestamp.............")
         element.ts
       }
     }
    )


    val result: DataStream[String] = dsWatermark.keyBy(_.word).timeWindow(Time.seconds(5))
      .process(new ProcessWindowFunction[WordBean, String, String, TimeWindow] {
        override def process(key: String, ctx: Context, elements: Iterable[WordBean], out: Collector[String]): Unit = {
          println("process......key: " + key + " List: " + elements.mkString("@"))
          out.collect("开始窗口：" + ctx.window.getStart + " 结束窗口：" + ctx.window.getEnd() + " 数据size：" + elements.size +
            " watermark: " + ctx.currentWatermark)
        }
      })
    result.print("result:")
    env.execute("Ascending")
  }
}
